/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.facebook.presto.adaptive.selection;

import com.facebook.airlift.log.Logger;
import com.facebook.presto.adaptive.AdaptiveMetadataManager;
import com.facebook.presto.adaptive.AdaptiveUtil;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.TableHandle;
import com.facebook.presto.spi.adaptive.CatalogEngine;
import com.facebook.presto.spi.adaptive.ConnectorAdaptiveOptimizer;
import com.facebook.presto.spi.adaptive.OptimizerType;
import com.facebook.presto.spi.adaptive.metadata.presto.AdaptivePrestoObject;
import com.facebook.presto.spi.plan.FilterNode;
import com.facebook.presto.spi.plan.PlanNode;
import com.facebook.presto.spi.plan.PlanNodeIdAllocator;
import com.facebook.presto.spi.plan.PlanVisitor;
import com.facebook.presto.spi.plan.TableScanNode;
import com.facebook.presto.spi.predicate.TupleDomain;
import com.facebook.presto.spi.relation.VariableReferenceExpression;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

import static java.util.Objects.requireNonNull;

/**
 * @author wenxu.Luo
 * UnionIndexSelector: a RBO adptive selector
 * <p>
 * metadata: in mongodb, column a and column b has union index.
 * query : select a,b from t1 where a>1 and b>2
 * result : => transform to mongodb.
 */
public class UnionIndexSelector
        extends RBOAdaptiveSelector
{

    private static final Logger log = Logger.get(UnionIndexSelector.class);

    // not all databases support booleans, so use 1=1 and 1=0 instead
    private static final String ALWAYS_TRUE = "1=1";
    private static final String ALWAYS_FALSE = "1=0";

    private final String quote = "\"";

    List<ConnectorAdaptiveOptimizer> connectorAdaptiveOptimizers;

    @Override
    public PlanNode adaptiveSelect(PlanNode maxSubplan,
            ConnectorSession session,
            PlanNodeIdAllocator idAllocator,
            List<ConnectorAdaptiveOptimizer> connectorAdaptiveOptimizers)
    {
        this.connectorAdaptiveOptimizers = connectorAdaptiveOptimizers;
        return maxSubplan.accept(new Visitor(session, idAllocator), null);
    }

    private class Visitor
            extends PlanVisitor<PlanNode, Void>
    {
        private final ConnectorSession session;
        private final PlanNodeIdAllocator idAllocator;

        public Visitor(ConnectorSession session, PlanNodeIdAllocator idAllocator)
        {
            this.session = requireNonNull(session, "session is null");
            this.idAllocator = requireNonNull(idAllocator, "idAllocator is null");
        }

        @Override
        public PlanNode visitPlan(PlanNode node, Void context)
        {
            ImmutableList.Builder<PlanNode> children = ImmutableList.builder();
            boolean changed = false;
            for (PlanNode child : node.getSources()) {
                PlanNode newChild = child.accept(this, null);
                if (newChild != child) {
                    changed = true;
                }
                children.add(newChild);
            }

            if (!changed) {
                return node;
            }
            return node.replaceChildren(children.build());
        }

        @Override
        public PlanNode visitTableScan(TableScanNode node, Void context)
        {

            TableHandle tableHandle = node.getTable();
            List<ColumnHandle> outputColumns = new ArrayList<>();

            // fill all columns and select columns with domain
            for (VariableReferenceExpression variable : node.getOutputVariables()) {
                ColumnHandle column = node.getAssignments().get(variable);
                outputColumns.add(column);
            }
            // select execute engine and optimize
            TableScanNode newNode = (TableScanNode) adaptiveSelectAndOptimize(node, tableHandle, outputColumns, new ArrayList<>(), idAllocator);

            return newNode;
        }

        @Override
        public PlanNode visitFilter(FilterNode node, Void context)
        {
            if (!(node.getSource() instanceof TableScanNode)) {
                return node;
            }
            TableScanNode tableScanNode = (TableScanNode) node.getSource();

            TableHandle tableHandle = ((TableScanNode) node.getSource()).getTable();
            TupleDomain<ColumnHandle> tupleDomain = tableHandle.getLayout().get().getTupleDomain();

            List<ColumnHandle> outputColumns = new ArrayList<>();
            List<ColumnHandle> domainColumns = new ArrayList<>();

            // fill all columns and select columns with domain
            for (VariableReferenceExpression variable : tableScanNode.getOutputVariables()) {
                ColumnHandle column = tableScanNode.getAssignments().get(variable);
                outputColumns.add(column);

                if (tupleDomain.getDomains().get().get(column) != null) {
                    domainColumns.add(column);
                }
            }

            // select execute engine and optimize
            FilterNode newNode = (FilterNode) adaptiveSelectAndOptimize(node, tableHandle, outputColumns, domainColumns, idAllocator);

            return newNode;
        }
    }

    /**
     * Key methods to select adaptive engine and optimize the maxSubTree.
     * 1. select the engine according rules in this class.
     * 2. check the best engine: (1) performance is better (2) engine has the backup of this table.
     * 3. send maxSubTree to selected engine to optimize.
     * 4. replace the old node.
     */
    private PlanNode adaptiveSelectAndOptimize(PlanNode node,
            TableHandle tableHandle,
            List<ColumnHandle> outputColumns,
            List<ColumnHandle> domainColumns,
            PlanNodeIdAllocator idAllocator)
    {
        CatalogEngine catalogEngine = selectPushDownEngine(tableHandle, outputColumns, domainColumns);
        if (catalogEngine == null) {
            return node;
        }

        PlanNode newNode = node;
        if (catalogEngine.equals(CatalogEngine.MONGODB)) {
            List<ConnectorAdaptiveOptimizer> optimizers = filterAdaptiveOptimizers(this.connectorAdaptiveOptimizers, CatalogEngine.MONGODB, OptimizerType.UNIONINDEX);
            if (optimizers == null) {
                return node;
            }

            AdaptivePrestoObject adaptiveObject = AdaptiveUtil.getMetadataManager()
                    .getAdaptivePrestoObject(CatalogEngine.MONGODB.toString(), null,
                            tableHandle.getConnectorHandle().getSchemaTableName().getTableName());

            for (ConnectorAdaptiveOptimizer optimizer : optimizers) {
                newNode = optimizer.adaptiveOptimize(newNode, tableHandle, idAllocator, adaptiveObject);
            }
        }
        return newNode;
    }

    /**
     * select push down engine.
     * two conditions:
     * 1. the engine has this table.
     * 2. the engine is satisfied push-down rules.
     */
    private CatalogEngine selectPushDownEngine(TableHandle tableHandle, List<ColumnHandle> outColumns, List<ColumnHandle> domainColumns)
    {
        String tableName = tableHandle.getConnectorHandle().getSchemaTableName().getTableName();

        AdaptiveMetadataManager metadataManager = AdaptiveUtil.getMetadataManager();
        List<String> candidateCatalog = metadataManager.getCatalogListByTableName(tableName);

        if (candidateCatalog.contains(CatalogEngine.MONGODB.toString()) && isSatisfiedMongoUnionIndex(tableName, outColumns, domainColumns)) {
            log.info("Adaptive selector - Select Engine, UnionIndexSelector selects mongodb engine.");
            return CatalogEngine.MONGODB;
        }
        return null;
    }

    /**
     * Check whether it's satisfied the rule.
     * In this rule, it requires output columns and filter columns must in the range of union indexes.
     * For example, table A has union index, column a and b.
     * select a,b,c from A  -- false
     * select a,b from A -- true
     * select a from A where a>10 and b>10 -- true
     * select a from A where c<10 -- false
     */
    private boolean isSatisfiedMongoUnionIndex(String tableName, List<ColumnHandle> outColumns, List<ColumnHandle> domainColumns)
    {
        AdaptiveMetadataManager metadataManager = AdaptiveUtil.getMetadataManager();
        AdaptivePrestoObject adaptiveObject = metadataManager
                .getAdaptivePrestoObject(CatalogEngine.MONGODB.toString(), null, tableName);
        List<String> unionIndexList = adaptiveObject.getPrestoTable().getMongoIndex();

        List<String> domColumns = new ArrayList<>();
        if (!domainColumns.isEmpty()) {
            domColumns = domainColumns.stream().map(ColumnHandle::getColumnName).collect(Collectors.toList());
        }

        List<String> outputColumns = outColumns.stream().map(ColumnHandle::getColumnName).collect(Collectors.toList());

        if (unionIndexList == null || unionIndexList.isEmpty()) {
            return false;
        }
        for (String unionIndex : unionIndexList) {
            List<String> indexList = Arrays.asList(unionIndex.split("\\|"));
            if (indexList.containsAll(domColumns) && indexList.containsAll(outputColumns)) {
                return true;
            }
        }

        return false;
    }


    /*---------------------------------------------test--------------------------------------------------*/

    @VisibleForTesting
    public PlanNode testScanAdaptiveSelect(TableScanNode node,
            List<String> indexList,
            PlanNodeIdAllocator idAllocator,
            List<ConnectorAdaptiveOptimizer> connectorAdaptiveOptimizers)
    {
        this.connectorAdaptiveOptimizers = connectorAdaptiveOptimizers;

        TableHandle tableHandle = node.getTable();
        List<ColumnHandle> outputColumns = new ArrayList<>();

        // fill all columns and select columns with domain
        for (VariableReferenceExpression variable : node.getOutputVariables()) {
            ColumnHandle column = node.getAssignments().get(variable);
            outputColumns.add(column);
        }

        return testAdaptiveSelectAndOptimize(node, tableHandle, outputColumns, new ArrayList<>(), idAllocator, indexList);
    }

    @VisibleForTesting
    public PlanNode testFilterAdaptiveSelect(FilterNode node,
            List<String> indexList,
            PlanNodeIdAllocator idAllocator,
            List<ConnectorAdaptiveOptimizer> connectorAdaptiveOptimizers)
    {
        this.connectorAdaptiveOptimizers = connectorAdaptiveOptimizers;

        if (!(node.getSource() instanceof TableScanNode)) {
            return node;
        }
        TableScanNode tableScanNode = (TableScanNode) node.getSource();

        TableHandle tableHandle = ((TableScanNode) node.getSource()).getTable();
        TupleDomain<ColumnHandle> tupleDomain = tableHandle.getLayout().get().getTupleDomain();

        List<ColumnHandle> outputColumns = new ArrayList<>();
        List<ColumnHandle> domainColumns = new ArrayList<>();

        // fill all columns and select columns with domain
        for (VariableReferenceExpression variable : tableScanNode.getOutputVariables()) {
            ColumnHandle column = tableScanNode.getAssignments().get(variable);
            outputColumns.add(column);

            if (tupleDomain.getDomains().get().get(column) != null) {
                domainColumns.add(column);
            }
        }

        return testAdaptiveSelectAndOptimize(node, tableHandle, outputColumns, domainColumns, idAllocator, indexList);
    }

    @VisibleForTesting
    private PlanNode testAdaptiveSelectAndOptimize(PlanNode node,
            TableHandle tableHandle,
            List<ColumnHandle> outputColumns,
            List<ColumnHandle> domainColumns,
            PlanNodeIdAllocator idAllocator,
            List<String> indexList)
    {
        CatalogEngine catalogEngine = testSelectPushDownEngine(tableHandle, outputColumns, domainColumns, indexList);
        if (catalogEngine == null) {
            return node;
        }

        PlanNode newNode = node;
        if (catalogEngine.equals(CatalogEngine.TEST)) {
            List<ConnectorAdaptiveOptimizer> optimizers = filterAdaptiveOptimizers(this.connectorAdaptiveOptimizers, CatalogEngine.TEST, OptimizerType.TEST);
            if (optimizers == null) {
                return node;
            }

            for (ConnectorAdaptiveOptimizer optimizer : optimizers) {
                newNode = optimizer.adaptiveOptimize(newNode, tableHandle, idAllocator, null);
            }
        }
        return newNode;
    }

    @VisibleForTesting
    private CatalogEngine testSelectPushDownEngine(TableHandle tableHandle, List<ColumnHandle> outColumns, List<ColumnHandle> domainColumns, List<String> indexList)
    {
        String tableName = tableHandle.getConnectorHandle().getSchemaTableName().getTableName();

        List<String> candidateCatalog = ImmutableList.of(CatalogEngine.TEST.toString());

        if (candidateCatalog.contains(CatalogEngine.TEST.toString()) && isSatisfiedTestUnionIndex(tableName, outColumns, domainColumns, indexList)) {
            log.info("Adaptive selector - Select Engine, UnionIndexSelector selects test engine.");
            return CatalogEngine.TEST;
        }
        return null;
    }

    @VisibleForTesting
    private boolean isSatisfiedTestUnionIndex(String tableName, List<ColumnHandle> outColumns, List<ColumnHandle> domainColumns, List<String> unionIndexList)
    {
        List<String> domColumns = new ArrayList<>();
        if (!domainColumns.isEmpty()) {
            domColumns = domainColumns.stream().map(ColumnHandle::getColumnName).collect(Collectors.toList());
        }

        List<String> outputColumns = outColumns.stream().map(ColumnHandle::getColumnName).collect(Collectors.toList());

        if (unionIndexList == null || unionIndexList.isEmpty()) {
            return false;
        }
        for (String unionIndex : unionIndexList) {
            List<String> indexList = Arrays.asList(unionIndex.split("\\|"));
            if (indexList.containsAll(domColumns) && indexList.containsAll(outputColumns)) {
                return true;
            }
        }

        return false;
    }


    /*---------------------------------------------test-finished--------------------------------------------------*/
}

